#!/usr/bin/env python3
"""
数据库结构更新脚本
用于将旧的数据库结构迁移到新的结构
"""

import os
import sys
import sqlite3
import json
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('db_migration.log')
    ]
)
logger = logging.getLogger("db-migration")

# 数据库路径
DB_PATH = os.environ.get('DB_PATH', os.path.join(os.path.dirname(os.path.abspath(__file__)), 'conversations.db'))

def backup_database(db_path):
    """备份数据库"""
    import shutil
    backup_path = f"{db_path}.bak"
    logger.info(f"备份数据库到 {backup_path}")
    shutil.copy2(db_path, backup_path)
    return backup_path

def update_schema(db_path):
    """更新数据库结构"""
    logger.info(f"开始更新数据库结构: {db_path}")
    
    with sqlite3.connect(db_path) as conn:
        # 检查是否已经有新字段
        cursor = conn.execute("PRAGMA table_info(conversations)")
        columns = [row[1] for row in cursor.fetchall()]
        
        # 如果已经有message_type字段，说明已经更新过
        if 'message_type' in columns:
            logger.info("数据库结构已经是最新的")
            return
        
        # 创建临时表
        logger.info("创建临时表")
        conn.execute("""
            CREATE TABLE conversations_new (
                id TEXT PRIMARY KEY,
                datetime TEXT NOT NULL,
                topic TEXT NOT NULL,
                keywords TEXT NOT NULL,
                summary TEXT NOT NULL,
                message_type TEXT NOT NULL,
                content TEXT NOT NULL,
                metadata TEXT,
                status TEXT NOT NULL DEFAULT 'active'
            )
        """)
        
        # 迁移数据
        logger.info("迁移数据到新表")
        conn.execute("""
            INSERT INTO conversations_new (
                id, datetime, topic, keywords, summary, 
                message_type, content, metadata, status
            )
            SELECT 
                id, datetime, topic, keywords, summary, 
                'user_message', content, '{}', status
            FROM conversations
        """)
        
        # 删除旧表
        logger.info("删除旧表")
        conn.execute("DROP TABLE conversations")
        
        # 重命名新表
        logger.info("重命名新表")
        conn.execute("ALTER TABLE conversations_new RENAME TO conversations")
        
        # 创建索引
        logger.info("创建索引")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_datetime ON conversations(datetime)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_status ON conversations(status)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_conversations_message_type ON conversations(message_type)")
        
        logger.info("数据库结构更新完成")

def clean_system_messages(db_path):
    """清理系统消息和临时数据"""
    logger.info("开始清理系统消息和临时数据")
    
    with sqlite3.connect(db_path) as conn:
        # 检查content字段中是否包含系统提示或思考过程
        cursor = conn.execute("SELECT id, content FROM conversations")
        rows = cursor.fetchall()
        
        cleaned_count = 0
        for row_id, content in rows:
            # 检查是否包含系统提示或思考过程的标记
            if "<thinking>" in content or "系统提示" in content or "system prompt" in content:
                # 将这些消息标记为系统消息
                conn.execute(
                    "UPDATE conversations SET message_type = 'system_message' WHERE id = ?",
                    (row_id,)
                )
                cleaned_count += 1
        
        logger.info(f"已标记 {cleaned_count} 条系统消息")

def main():
    """主函数"""
    try:
        # 备份数据库
        backup_path = backup_database(DB_PATH)
        logger.info(f"数据库已备份到: {backup_path}")
        
        # 更新数据库结构
        update_schema(DB_PATH)
        
        # 清理系统消息
        clean_system_messages(DB_PATH)
        
        logger.info("数据库迁移完成")
    except Exception as e:
        logger.error(f"数据库迁移失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()